from rest_framework.renderers import JSONRenderer


# 控制返回的JSON格式的类
class CustomRenderer(JSONRenderer):
    def render(self, data, accepted_media_type=None, renderer_context=None):

        if renderer_context:
            msg = '成功'
            code = renderer_context['response'].status_code
            # 判断实例的类型，返回的数据可能是列表也可能是字典
            if isinstance(data, dict):
                # 如果是字典的话应该是返回的数据，会包含 msg, code, status 等字段，必须抽离出来
                msg = data.pop('msg', '成功')

                if data.get('code', None):
                    code = data.pop('code')

                # 把data 抽离出来
                if data.get('data', None):
                    data = data.pop('data')


            # 自定义返回数据格式
            res = {
                'msg': msg,
                'code': code,
                'data': data
            }

            # 返回 JSON 数据
            return super().render(res, accepted_media_type, renderer_context)
        else:
            return super().render(data, accepted_media_type, renderer_context)
